#include "efficient_sam.h"

// Hash function for images
std::size_t compute_image_hash(const cv::Mat& img)
{
	std::size_t hash = 0;
	for (int i = 0; i < img.total() * img.elemSize(); ++i)
	{
		hash = hash * 101 + img.data[i];
	}
	return hash;
}

EfficientSam::EfficientSam(const std::string& encoder_path, const std::string& decoder_path)
	: env(ORT_LOGGING_LEVEL_WARNING, "EfficientSam"), session_options(),
	  encoder_session(nullptr), decoder_session(nullptr)
{
	session_options.SetIntraOpNumThreads(1);

	// Convert std::string to std::wstring
	std::wstring wencoder_path(encoder_path.begin(), encoder_path.end());
	std::wstring wdecoder_path(decoder_path.begin(), decoder_path.end());

	// Pass wchar_t* to Ort::Session constructor
	encoder_session = Ort::Session(env, wencoder_path.c_str(), session_options);
	decoder_session = Ort::Session(env, wdecoder_path.c_str(), session_options);
}

void EfficientSam::set_image(const cv::Mat& image_input)
{
	std::lock_guard<std::mutex> lock(mtx);
	image = image_input.clone();
	std::size_t image_hash = compute_image_hash(image);

	auto it = image_embedding_cache.find(image_hash);
	if (it != image_embedding_cache.end())
	{
		image_embedding = it->second;
	}
	else
	{
		if (worker_thread.joinable())
		{
			worker_thread.join();
		}
		worker_thread = std::thread(&EfficientSam::compute_and_cache_image_embedding, this, image_hash);
	}
}

void EfficientSam::compute_and_cache_image_embedding(std::size_t image_hash)
{
	std::lock_guard<std::mutex> lock(mtx);
	std::cout << "Computing image embedding..." << std::endl;

	cv::Mat rgb_image;
	if (image.channels() == 4)
	{
		cv::cvtColor(image, rgb_image, cv::COLOR_RGBA2RGB);
	}
	else
	{
		rgb_image = image.clone();
	}

	// Convert image to float and normalize
	cv::Mat float_image;
	rgb_image.convertTo(float_image, CV_32FC3, 1.0 / 255.0);

	// Transpose to (1, 3, H, W)
	std::vector<cv::Mat> channels(3);
	cv::split(float_image, channels);
	std::vector<float> input_tensor_values;
	for (const auto& ch : channels)
	{
		input_tensor_values.insert(input_tensor_values.end(), (float*)ch.datastart, (float*)ch.dataend);
	}

	std::vector<int64_t> input_shape = {1, 3, rgb_image.rows, rgb_image.cols};

	// Create input tensor
	Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
	Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
		memory_info, input_tensor_values.data(), input_tensor_values.size(),
		input_shape.data(), input_shape.size());

	// Run encoder session
	const char* input_names[] = {"batched_images"};
	auto output_tensors = encoder_session.Run(
		Ort::RunOptions{nullptr}, input_names, &input_tensor, 1, nullptr, 0);

	// Retrieve output
	Ort::Value& output_tensor = output_tensors.front();
	float* output_data = output_tensor.GetTensorMutableData<float>();
	std::vector<int64_t> output_shape = output_tensor.GetTensorTypeAndShapeInfo().GetShape();

	std::vector<int> output_shape_int(output_shape.begin(), output_shape.end());
	// Create image_embedding from output data
	image_embedding = cv::Mat(output_shape_int, CV_32F, output_data).clone();

	// Cache the embedding
	if (image_embedding_cache.size() > 10)
	{
		image_embedding_cache.erase(image_embedding_cache.begin());
	}
	image_embedding_cache[image_hash] = image_embedding;

	std::cout << "Done computing image embedding." << std::endl;
}

cv::Mat EfficientSam::get_image_embedding()
{
	if (worker_thread.joinable())
	{
		worker_thread.join();
	}
	std::lock_guard<std::mutex> lock(mtx);
	return image_embedding;
}

cv::Mat EfficientSam::predict_mask_from_points(const std::vector<cv::Point2f>& points,
                                               const std::vector<float>& point_labels)
{
	cv::Mat image_embedding = get_image_embedding();

	// Prepare point coordinates and labels
	std::vector<float> input_point_values;
	for (const auto& pt : points)
	{
		input_point_values.push_back(pt.x);
		input_point_values.push_back(pt.y);
	}

	std::vector<int64_t> point_coords_shape = {1, 1, (int64_t)points.size(), 2};
	std::vector<int64_t> point_labels_shape = {1, 1, (int64_t)point_labels.size()};

	// Prepare orig_im_size
	std::vector<int64_t> orig_im_size = {image.rows, image.cols};

	// Create tensors
	Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
	Ort::Value batched_point_coords_tensor = Ort::Value::CreateTensor<float>(
		memory_info, input_point_values.data(), input_point_values.size(),
		point_coords_shape.data(), point_coords_shape.size());

	float* point_labels_data = const_cast<float*>(point_labels.data());
	Ort::Value batched_point_labels_tensor = Ort::Value::CreateTensor<float>(
		memory_info, point_labels_data, point_labels.size(),
		point_labels_shape.data(), point_labels_shape.size());

	Ort::Value orig_im_size_tensor = Ort::Value::CreateTensor<int64_t>(
		memory_info, orig_im_size.data(), orig_im_size.size(),
		std::vector<int64_t>{2}.data(), 1);

	// Prepare inputs
	const char* input_names[] = {"image_embeddings", "batched_point_coords", "batched_point_labels", "orig_im_size"};
	float* image_embedding_data = image_embedding.ptr<float>();
	std::vector<int64_t> embedding_shape(image_embedding.size.p, image_embedding.size.p + image_embedding.dims);
	Ort::Value image_embeddings_tensor = Ort::Value::CreateTensor<float>(
		memory_info, image_embedding_data, image_embedding.total(),
		embedding_shape.data(), embedding_shape.size());

	std::array<Ort::Value, 4> input_tensors = {
		std::move(image_embeddings_tensor),
		std::move(batched_point_coords_tensor),
		std::move(batched_point_labels_tensor),
		std::move(orig_im_size_tensor)
	};

	// Run decoder session
	auto output_tensors = decoder_session.Run(
		Ort::RunOptions{nullptr}, input_names, input_tensors.data(), input_tensors.size(), nullptr, 0);

	// Process output
	Ort::Value& masks_tensor = output_tensors.front();
	float* masks_data = masks_tensor.GetTensorMutableData<float>();
	std::vector<int64_t> masks_shape = masks_tensor.GetTensorTypeAndShapeInfo().GetShape();

	int H = masks_shape.back() - 2;
	int W = masks_shape.back() - 1;
	cv::Mat mask(H, W, CV_32F, masks_data);
	cv::Mat binary_mask;
	cv::threshold(mask, binary_mask, 0.0, 1.0, cv::THRESH_BINARY);

	// Remove small objects
	double min_size_ratio = 0.05;
	int min_size = static_cast<int>(cv::countNonZero(binary_mask) * min_size_ratio);
	cv::Mat labels, stats, centroids;
	int num_labels = cv::connectedComponentsWithStats(binary_mask, labels, stats, centroids);

	for (int i = 1; i < num_labels; ++i)
	{
		int area = stats.at<int>(i, cv::CC_STAT_AREA);
		if (area < min_size)
		{
			binary_mask.setTo(0, labels == i);
		}
	}

	return binary_mask;
}

cv::Mat EfficientSam::predict_polygon_from_points(const std::vector<cv::Point2f>& points,
                                                  const std::vector<float>& point_labels)
{
	cv::Mat mask = predict_mask_from_points(points, point_labels);

	// Implement compute_polygon_from_mask
	std::vector<std::vector<cv::Point>> contours;
	cv::findContours(mask, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

	cv::Mat polygon = cv::Mat::zeros(mask.size(), CV_8UC1);
	cv::drawContours(polygon, contours, -1, cv::Scalar(255), cv::FILLED);

	return polygon;
}
